// CUDA kernels and wrappers for GPU-bound image processing

#include <stdlib.h>
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>

__device__ __forceinline__ unsigned char Clamp(int val)
{
    val = val > 255 ? 255 : val;
    return val < 0 ? 0 : val;
}

int* MapIndices;

__global__ void YUVtoBGRA(int* src, long long* dst, int width)
{
	int xIndex = (blockIdx.x * blockDim.x + threadIdx.x) * 2;
	int yIndex = blockIdx.y * blockDim.y + threadIdx.y;  
	int index = (yIndex*width + xIndex)/2;

	int val = src[index];
	unsigned char* ptr = reinterpret_cast<unsigned char*>(&val);
	long long out = 0;
	unsigned char* outPtr = reinterpret_cast<unsigned char*>(&out);

	int c1 = ((ptr[1] - 16) * 298) >> 8;
	int c2 = ((ptr[3] - 16) * 298) >> 8;
	int d = ptr[0] - 128;
	int e = ptr[2] - 128;

	int r_add = (409 * e + 128) >> 8;
	int g_add = (-100 * d - 208 * e + 128) >> 8;
	int b_add = (516 * d + 128) >> 8;

	outPtr[2] = Clamp(c1 + r_add);
	outPtr[1] = Clamp(c1 + g_add);
	outPtr[0] = Clamp(c1 + b_add);
	outPtr[3] = 255;

	outPtr[6] = Clamp(c2 + r_add);
	outPtr[5] = Clamp(c2 + g_add);
	outPtr[4] = Clamp(c2 + b_add);
	outPtr[7] = 255;

	dst[index] = out;
}

// Aligns the color frame to the depth frame based on the alignment map
__global__ void AlignColorFrame(int* src, int* dst, int* mapIndices)
{
    unsigned int index = blockIdx.x*blockDim.x + threadIdx.x;
	dst[index] = src[mapIndices[index]];
}

// Zeros all pixels in the depth frame that don't belong to the registered player
__global__ void BackgroundSubtract(short* depthFrame, int playerIndex)
{
    int index = blockIdx.x*blockDim.x + threadIdx.x;
	short depth = depthFrame[index];
	if((depth & 0x07) == playerIndex)
		depthFrame[index] = depth >> 3;
	else
		depthFrame[index] = 0;
}

// Zeros all pixels in the depth frame that don't belong to the registered player
__global__ void BackgroundSubtractWithAlpha(short* depthFrame, unsigned char* alphaChannel, int playerIndex)
{
    unsigned int index = blockIdx.x*blockDim.x + threadIdx.x;
	if((depthFrame[index] & 0x07) == playerIndex)
	{
		depthFrame[index] >>= 3;
		alphaChannel[index] = 0xFF;
	}
	else
	{
		depthFrame[index] = 0;
		alphaChannel[index] = 0;
	}
}

// Wrapper for YUV to BGRA conversion
extern "C"
cudaError_t CUDA_YUVtoBGRA(unsigned char* src, unsigned char* dst, int width, int height)
{
	dim3 numBlocks(width / 32, height / 16, 1);
	dim3 numThreadsPerBlock(16, 16, 1);

	YUVtoBGRA<<<numBlocks, numThreadsPerBlock>>>(
		reinterpret_cast<int*>(src),
		reinterpret_cast<long long*>(dst),
		width);

	return cudaGetLastError();
};

// Wrapper for color frame alignment
extern "C"
cudaError_t CUDA_AlignColorFrame(int* src, int* dst, int length)
{
	dim3 numBlocks(length / 256, 1, 1);
	dim3 numThreadsPerBlock(256, 1, 1);
	AlignColorFrame<<<numBlocks, numThreadsPerBlock>>>(src, dst, MapIndices);
	cudaDeviceSynchronize();

	return cudaGetLastError();
}

// Wrapper for background subtraction
extern "C"
cudaError_t CUDA_BackgroundSubtract(unsigned char* depthFrame, unsigned char* alphaChannel, int length, int playerIndex)
{
	dim3 numBlocks(length / 256, 1, 1);
	dim3 numThreadsPerBlock(256, 1, 1);
	if(alphaChannel == NULL)
		BackgroundSubtract<<<numBlocks, numThreadsPerBlock>>>(reinterpret_cast<short*>(depthFrame), playerIndex);
	else
		BackgroundSubtractWithAlpha<<<numBlocks, numThreadsPerBlock>>>(reinterpret_cast<short*>(depthFrame), alphaChannel, playerIndex);
	
	return cudaGetLastError();
};

extern "C"
cudaError_t CUDA_SetMapIndices(int* mapIndices)
{
	cudaError_t error = cudaSuccess;
	if(MapIndices != NULL)
		error = cudaFree(MapIndices);
	MapIndices = mapIndices;
	return error;
}